#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/Xatom.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/time.h>
#include <png.h>

int get_offset(long mask)
{
    int n = 0;
    while(mask && (mask & 0x1) == 0)
    {
        n++;
        mask >>= 1;
    }
    return n;
}
// 保存为 PPM 格式（简单图像格式）
void save_as_ppm(XImage *image, const char *filename) {
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        perror("无法打开文件");
        return;
    }
    
    // PPM 文件头
    fprintf(fp, "P6\n%d %d\n255\n", image->width, image->height);
    
    int red_offset = get_offset(image->red_mask);
    int green_offset = get_offset(image->green_mask);
    int blue_offset = get_offset(image->blue_mask);
    // 写入像素数据
    for (int y = 0; y < image->height; y++) {
        for (int x = 0; x < image->width; x++) {
            unsigned long pixel = XGetPixel(image, x, y);
            
            // 提取 RGB 分量
            unsigned char r = (pixel & image->red_mask) >> red_offset;
            unsigned char g = (pixel & image->green_mask) >> green_offset;
            unsigned char b = (pixel & image->blue_mask) >> blue_offset;
            
            // 写入文件
            fputc(r, fp);
            fputc(g, fp);
            fputc(b, fp);
        }
    }
    
    fclose(fp);
    printf("截图已保存为: %s\n", filename);
}

// 保存为 PNG 格式（需要 libpng）
void save_as_png(XImage *image, const char *filename) {
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        perror("无法打开PNG文件");
        return;
    }
    
    png_structp png_ptr = png_create_write_struct(
        PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!png_ptr) {
        fclose(fp);
        return;
    }
    
    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_write_struct(&png_ptr, NULL);
        fclose(fp);
        return;
    }
    
    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_write_struct(&png_ptr, &info_ptr);
        fclose(fp);
        return;
    }
    
    png_init_io(png_ptr, fp);
    
    // 设置PNG头部
    png_set_IHDR(
        png_ptr, info_ptr,
        image->width, image->height,
        8, // 位深度
        PNG_COLOR_TYPE_RGB,
        PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_DEFAULT,
        PNG_FILTER_TYPE_DEFAULT
    );
    
    png_write_info(png_ptr, info_ptr);
    
    // 分配行缓冲区
    png_bytep *row_pointers = (png_bytep*)malloc(sizeof(png_bytep) * image->height);
    for (int y = 0; y < image->height; y++) {
        row_pointers[y] = (png_byte*)malloc(png_get_rowbytes(png_ptr, info_ptr));
    }
    int red_offset = get_offset(image->red_mask);
    int green_offset = get_offset(image->green_mask);
    int blue_offset = get_offset(image->blue_mask);
    
    // 转换像素数据
    for (int y = 0; y < image->height; y++) {
        png_bytep row = row_pointers[y];
        for (int x = 0; x < image->width; x++) {
            unsigned long pixel = XGetPixel(image, x, y);
            
            // 提取RGB分量
            *row++ = (pixel & image->red_mask) >> red_offset;   // R
            *row++ = (pixel & image->green_mask) >> green_offset; // G
            *row++ = (pixel & image->blue_mask) >> blue_offset;  // B
        }
    }
    
    // 写入图像数据
    png_write_image(png_ptr, row_pointers);
    png_write_end(png_ptr, NULL);
    
    // 清理资源
    for (int y = 0; y < image->height; y++) {
        free(row_pointers[y]);
    }
    free(row_pointers);
    
    png_destroy_write_struct(&png_ptr, &info_ptr);
    fclose(fp);
    printf("截图已保存为PNG: %s\n", filename);
}

// 获取窗口截图
XImage* capture_window(Display *display, Window window, int x, int y, unsigned int width, unsigned int height) {
    // 获取窗口图像
    XImage *image = XGetImage(
        display, 
        window, 
        0, 0,             // 起始坐标
        width, height,    // 尺寸
        AllPlanes,        // 所有平面
        ZPixmap           // 格式
    );
    printf("image: depth:%d, bytes_per_line:%d. bits_per_pixel:%d, " 
           "red_mask:%lx, green_mask:%lx, blue_mask:%lx\n", 
           image->depth, image->bytes_per_line, image->bits_per_pixel,
           image->red_mask, image->green_mask, image->blue_mask);
    
    return image;
}

// 获取活动窗口
Window get_active_window(Display *display) {
    Atom net_active_window = XInternAtom(display, "_NET_ACTIVE_WINDOW", False);
    Atom actual_type;
    int actual_format;
    unsigned long nitems, bytes_after;
    unsigned char *data = NULL;
    Window active_window = 0;
    
    if (XGetWindowProperty(
        display, 
        DefaultRootWindow(display), 
        net_active_window, 
        0, 1, 
        False, 
        XA_WINDOW, 
        &actual_type, 
        &actual_format, 
        &nitems, 
        &bytes_after, 
        &data) == Success && data) {
        
        if (nitems > 0) {
            active_window = *(Window*)data;
        }
        XFree(data);
    }
    
    return active_window;
}

// 获取窗口几何信息
void get_window_geometry(Display *display, Window window, 
                         int *x, int *y, 
                         unsigned int *width, unsigned int *height) {
    Window root;
    int orig_x, orig_y;
    unsigned int border_width, depth;
    
    XGetGeometry(display, window, &root, 
                &orig_x, &orig_y, 
                width, height, 
                &border_width, &depth);
    
    XTranslateCoordinates(display, window, root, 
                          0, 0, 
                          x, y, 
                          &root);
}

int main(int argc, char *argv[]) {
    Display *display = XOpenDisplay(NULL);
    if (!display) {
        fprintf(stderr, "无法打开X显示\n");
        return 1;
    }
    
    Window target_window = None;
    int x = 0, y = 0;
    unsigned int width = 0, height = 0;
    
    // 解析命令行参数
    if (argc > 1) {
        // 从参数获取窗口ID
        target_window = strtoul(argv[1], NULL, 0);
        
        if (argc >= 6) {
            // 获取自定义区域
            x = atoi(argv[2]);
            y = atoi(argv[3]);
            width = atoi(argv[4]);
            height = atoi(argv[5]);
        }
    } else {
        // 获取活动窗口
        target_window = get_active_window(display);
        if (target_window == None) {
            fprintf(stderr, "无法获取活动窗口\n");
            XCloseDisplay(display);
            return 1;
        }
    }
    
    // 获取窗口几何信息（如果未指定区域）
    if (width == 0 || height == 0) {
        get_window_geometry(display, target_window, &x, &y, &width, &height);
    }
    
    printf("截图窗口: 0x%lx\n", target_window);
    printf("位置: (%d, %d) 尺寸: %ux%u\n", x, y, width, height);
    
    // 捕获窗口图像
    XImage *image = capture_window(display, target_window, x, y, width, height);
    if (!image) {
        fprintf(stderr, "无法获取窗口图像\n");
        XCloseDisplay(display);
        return 1;
    }
    
    // 生成文件名
    char filename[256];
    struct timeval tv;
    gettimeofday(&tv, NULL);
    snprintf(filename, sizeof(filename), "screenshot_%ld.ppm", tv.tv_sec);
    
    // 保存图像
    save_as_ppm(image, filename);
    gettimeofday(&tv, NULL);
    snprintf(filename, sizeof(filename), "screenshot_%ld.png", tv.tv_sec);
    save_as_png(image, filename);
    
    // 清理资源
    XDestroyImage(image);
    XCloseDisplay(display);
    
    return 0;
}
